# pre packages
from myglobal import *

# sys packages
from pylab import np
import obspy as ob
from glob import glob
from tqdm import tqdm
from obspy.core.utcdatetime import UTCDateTime
import pandas as pd
from scipy import signal


# self packages
from utils.loc import load_loc, get_distance,sort_data_by_distance
from utils.h5data import h5glob, save_h5_data, read_h5_data
from utils.hsr import RawDataBase, search_events,hsr_vel_detection
from utils.plot import plot_traces_by_subfigures, plot_events
from utils.trace import get_tapered_slices, get_traces_attributes,get_traces_sqrtmean

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-debug', action='store_true', help='method: DEBUG')
parser.add_argument('-input', default='', help='1.py中生成的h5文件')
parser.add_argument('-line', default='', help='1.py中生成的h5文件')
parser.add_argument('-amp', default=-1,type=float, help='高铁信号幅度阈值')

args = parser.parse_args()
print(args)
DEBUG= args.debug
DATA_FILE = args.input



groups = read_h5_data(DATA_FILE,keys=['all_groups'], FILL_NONE=True, FILL_VALUE=None)[0]
if groups is None:
    print('search groups')
    groups = h5glob(DATA_FILE,pattern='DAY*/IDX*.T*',object_type='group')
    groups.sort()
    save_h5_data(DATA_FILE,{'all_groups':np.array(groups,dtype='S')}, mode='a')
else:
    groups = [g.decode() for g in groups]
    groups.sort()

_,args_infile = read_h5_data(DATA_FILE,keys=[], group_name='metadata',read_attrs=True)
LINE_ID = args_infile['line'] if args.line=='' else args.line
BASE_STAT=args_infile['base']
AMP_HSR = args_infile['amp'] if args.amp<=0 else args.amp
DATE = args_infile['date']
EVENT_L = args_infile['L']

date_info = get_info(DATE)
days = date_info['days']
H5_ROOT=date_info['H5_ROOT']
RAW_ROOT = date_info['RAW_ROOT']
s_info=date_info['s_info']
LINES = date_info['LINES']

if LINE_ID in LINES.keys():
    STATS = LINES[LINE_ID]
else:
    STATS = [LINE_ID]

FIG_ROOT = f'./figures/2.eventsInfo/{BASE_STAT}.Line{LINE_ID}'
if not os.path.exists(FIG_ROOT):
    os.makedirs(FIG_ROOT)

nx = len(STATS)
ne = len(groups)
 
AMP = []
AMP_HEALTH = []
VINFO = np.zeros([ne,7])
N_FIG = 100

IT = tqdm(enumerate(groups),desc="processes event data",total=ne)

for i,group_i in IT:
    
    data_infile, t, e_time, stats_infile= read_h5_data(DATA_FILE, keys=['data','t','te','stats'],group_name=group_i)
    stats_infile = [i.decode() for i in stats_infile]
    e_time = e_time.decode()
    dt = t[1]-t[0]
    nt = len(t)

    x = np.zeros([nx])
    data_i = np.zeros([nx,nt],dtype=np.float32)
    for j,name in enumerate(STATS):
        idx_j = stats_infile.index(name)
        xj, _,_ = get_distance(s_info=s_info,name1=BASE_STAT,name2=name, S2N=True)
        x[j]=xj
        data_i[j,:] = data_infile[idx_j,:]

    N_H = 1/dt

    # E_i,t_health = get_traces_attributes(np.array(traces), t, 2, 1, ATTR = 'E')
    # health_i = np.ones_like(E_i)
    # health_i[E_i>=AMP_HSR/3] = 0
    # AMP_HEALTH.append(health_i)
    # AMP.append(E_i)

    E_i,t_health = get_traces_sqrtmean(np.array(data_i), t, WIN_L=0.5)
    health_i = np.ones_like(E_i)
    health_i[E_i>=AMP_HSR*3] = 0
    
    AMP_HEALTH.append(health_i)
    AMP.append(E_i)

    # 检测事件速度
    traces4plot = np.diff(1-health_i,n=1, axis=-1)
    t4plot = (t_health[1:]+t_health[:-1])/2
    x4plot = x 
    EV_MIN = nx/10
    print(AMP_HSR, EV_MIN,E_i.max())
    detected_events_PV = hsr_vel_detection(traces4plot, x4plot, t4plot,
                                                          velocity_range=[60,120],
                                                          EMIN=EV_MIN)
    detected_events_NV = hsr_vel_detection(traces4plot, x4plot, t4plot,
                                                         velocity_range=[-120,-60],
                                                         EMIN=EV_MIN)

    VINFO[i,1:4]=detected_events_PV[0]
    VINFO[i,4:7]=detected_events_NV[0]
    
    # S2N
    if  detected_events_PV[0][-1] >=EV_MIN and detected_events_NV[0][-1]<EV_MIN:
        VINFO[i,0]=1
    # N2S
    elif detected_events_PV[0][-1]<EV_MIN and detected_events_NV[0][-1]>=EV_MIN:
        VINFO[i,0]=2
    # 双向
    elif detected_events_PV[0][-1]>=EV_MIN and detected_events_NV[0][-1]>=EV_MIN:
        VINFO[i,0]=3
    # 无信号
    else:
        VINFO[i,0]=0

    if DEBUG and i%int(ne/N_FIG)==0:
    
        fig, ax = plot_events(traces4plot, x4plot, t4plot , DO_FILTER=False, PLOT_WIGGLE=True, NORM=False, SCALE=400)
    
        for v,a,e in detected_events_PV:
            if e<EV_MIN:
                continue
            ax.plot(t4plot,(t4plot+a)*v, lw=1, color='blue')
            
        for v,a,e in detected_events_NV:
            if e<EV_MIN:
                continue
            ax.plot(t4plot,(t4plot+a)*v, lw=1, color='red')

        ax.set_title(f'{e_time},{VINFO[i]}')
        fig.tight_layout()
        filename = f'{FIG_ROOT}/{i:04d}.{e_time}.png'
        print(filename)
        fig.savefig(filename)

AMP_HEALTH = np.array(AMP_HEALTH)
AMP = np.array(AMP)

save_h5_data(DATA_FILE,{'all_groups':np.array(groups,dtype='S'),
                        'stats':np.array(STATS,dtype='S'),
                        'x':x,
                        'HEALTH':AMP_HEALTH,
                        'AMP':AMP,
                        'VINFO':VINFO,
                        't_health':t_health},
                        group_name='info', mode='a')